classdef BlochEquation < handle
    %BlochEquation Summary of this class goes here
    %   Detailed explanation goes here
    
    properties
        t
        domain
        kernel
        rhs

        range
        T1
        T2
        
        R
        omegaX
        omegaY
        omegaZ
        
        period
        initialM
        steadyM
        solution
        
        sol_opt
        steady_opt
    end
    
    methods
        function obj = BlochEquation(T1, T2, varargin)
            p = inputParser;
            p.KeepUnmatched = true;
            
            rgTest = @(x) length(x) == 2 && x(2)>x(1);
            t1Test = @(x) x > 0;
            t2Test = @(x) x > 0;
            
            stateTest = @(x) length(x) == 3 && isnumeric(x);
            sTest = @(x) any(ismember({'transient', 'periodic'}, x));
            tTest = @(x) any(ismember({'automatic', 'manual'}, x));
                        
            p.addRequired('T1',     t1Test);
            p.addRequired('T2',     t2Test);
            
            p.addParameter('tRange',  [0 5.0*max([T1, T2])], rgTest);
            p.addParameter('period',  inf,               @isnumeric);
            p.addParameter('init',    [0 0 0].',  stateTest);
            p.addParameter('steady',  [0 0 0].',  stateTest);
            p.addParameter('solution',  'transient',  sTest);
            p.addParameter('steadyOpt', 'automatic',  tTest);
            
            p.parse(T1, T2, varargin{:});
            
            obj.T1 = p.Results.T1;
            obj.T2 = p.Results.T2;
            
            obj.range = p.Results.tRange;
            obj.period   = p.Results.period;
            obj.initialM = p.Results.init;
            obj.steadyM  = p.Results.steady;
            obj.sol_opt  = p.Results.solution;
            obj.steady_opt = p.Results.steadyOpt;

            switch obj.sol_opt
                case 'transient'
                    obj.domain = obj.range;
                    if any(ismember(p.UsingDefaults, 'tRange'))
                        warning('Using default Range [0 5.0*max([T1, T2])] = [0 %f].', obj.range(2) );
                    end
                case 'periodic'
                    obj.domain = [0, obj.period];
                    if any(ismember(p.UsingDefaults, 'period'))
                        warning('Using default period [0 Inf].');
                    end
                otherwise
                    error('wrong solution option: %s', obj.sol_opt);
            end

            if strcmp(obj.steady_opt, 'manual') && any(ismember(p.UsingDefaults, 'steady'))
                warning('Using default steady state [0 0 0].;');
            end
            
            obj.t = chebfun('t', obj.domain);
        end
        
        function obj = set_pumping(obj, R)
            p = inputParser;
            chebTest = @(x) isnumeric(x) || isa(x, 'chebfun');
            p.addRequired('R', chebTest);
            p.parse(R);
           
            obj.R  = p.Results.R;
        end
        
        function obj = set_Omega(obj, varargin)
            p = inputParser;
            chebTest = @(x) isnumeric(x) || isa(x, 'chebfun');
            p.addParameter('omegaX',  obj.omegaX,      chebTest);
            p.addParameter('omegaY',  obj.omegaY,      chebTest);
            p.addParameter('omegaZ',  obj.omegaZ,      chebTest);
            p.parse(varargin{:});
            
            obj.omegaX   = p.Results.omegaX;
            obj.omegaY   = p.Results.omegaY;
            obj.omegaZ   = p.Results.omegaZ;
        end
        
        function obj = make_kernel(obj)
            ker = chebop(min(obj.domain), max(obj.domain));

            gt = 1/obj.T1+1/obj.T2+obj.R; 
            g1 = 1/obj.T1+obj.R;

            ker.op=@(t, mx, my, mz) [diff(mx) +         gt*mx + obj.omegaZ*my - obj.omegaY*mz; ...
                                     diff(my) - obj.omegaZ*mx +         gt*my + obj.omegaX*mz; ...
                                     diff(mz) + obj.omegaY*mx - obj.omegaX*my +         g1*mz];
                               
            switch obj.sol_opt
                case 'transient'
                    ker.lbc = @(mx, my, mz) [mx - obj.initialM(1); ...
                                             my - obj.initialM(2); ...
                                             mz - obj.initialM(3)];
                case 'periodic'
                    ker.bc = 'periodic';
                otherwise
                    error('wrong sol_opt: %s', obj.sol_opt);
            end
            
            obj.kernel = ker;
        end
        
        function obj = make_rhs(obj)
            switch obj.steady_opt
                case 'automatic'
                    Minf = [0.5* obj.R/( obj.R + 1/obj.T1); 0; 0];
                    rhs_op = [obj.R*Minf(1); obj.R*Minf(2); obj.R*Minf(3)];
                case 'manual'
                    Minf = obj.steadyM;
                    rhs_op = Minf;
                otherwise
                    error('wrong steady option: %s', obj.steady_opt);
            end
            
            obj.rhs = rhs_op;
        end
        
        function obj = solve(obj)
            tic;
            obj.make_kernel().make_rhs();
            sol = obj.kernel\obj.rhs;
            obj.solution = sol;
            toc;
        end
        
        function sol = getSolution(obj)
            sol.t = obj.t;
            sol.mx = obj.solution{1};
            sol.my = obj.solution{2};
            sol.mz = obj.solution{3};
        end
        
        function plot_solution(obj, sol)
            if nargin == 1
                sol = obj.getSolution();
            end
            
            c=colormap('lines');
            
            subplot(3, 2, [1 2])
            plot(sol.mx, 'b.-', sol.my, 'r.-');
            grid minor;
            
            subplot(3, 2, 3);
            plot3(sol.mx, sol.my, sol.mz, 'k.-')
            grid minor;
            ax = gca;
            ax.XColor = c(1, :);
            ax.YColor = c(2, :);
            ax.ZColor = c(3, :);
            ax.Box = 'off';
            xlabel('mx'); ylabel('my'); zlabel('mz');

            subplot(3, 2, 4);
            plot(sol.mx, sol.my, 'k.-')
            axis equal; grid minor;
            ax = gca;
            ax.XColor = c(1, :);
            ax.YColor = c(2, :);
            ax.Box = 'off';
            xlabel('m_x'); ylabel('m_y'); 

            subplot(3, 2, 5);
            plot(sol.my, sol.mz, 'k.-')
            axis equal; grid minor;
            ax = gca;
            ax.XColor = c(2, :);
            ax.YColor = c(3, :);
            ax.Box = 'off';
            xlabel('m_y'); ylabel('m_z'); 
            
            subplot(3, 2, 6);
            plot(sol.mx, sol.mz, 'k.-')
            ax = gca;
            axis equal; grid minor;
            ax.XColor = c(1, :);
            ax.YColor = c(3, :);
            ax.Box = 'off';
            xlabel('m_x'); ylabel('m_z'); 

        end
        
        function demod = harmonic(obj, component, order)
            if ~strcmp(obj.sol_opt, 'periodic')
                error('solution option must be periodic');
            end
            
            if nargin == 2
                order = 1;
            end
            
            sol = obj.getSolution();
            
            switch component
                case {'x', 'X', 1}
                    s = sol.mx;
                case {'y', 'Y', 2}
                    s = sol.my;
                case {'z', 'Z', 3}
                    s = sol.mz;
            end                    
            
            k = 1;
            demod=struct;
            for m = order
                demod(k).X = s'*cos(2*pi*m/obj.period*obj.t) * 2/obj.period;
                demod(k).Y = s'*sin(2*pi*m/obj.period*obj.t) * 2/obj.period;
                demod(k).R = sqrt(demod(k).X*demod(k).X+demod(k).Y*demod(k).Y);
                demod(k).phase = atan(demod(k).X/demod(k).Y);
                k = k + 1;
            end
        end
        
    end
    
    methods
        R = cosine_pump(obj, amp, freq, offset)
        R = square_pump(obj, amp, freq, offset)
    end
end

